import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.MessageDigest;
import java.util.Arrays;
import java.util.Scanner;

public class FileFolderEncryptor {

    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/ECB/PKCS5Padding";
    // 加密文件标记（用于区分加密/未加密文件）
    private static final String ENCRYPT_FLAG = ".encrypted";

    private static SecretKeySpec generateKey(String password) throws Exception {
        MessageDigest sha = MessageDigest.getInstance("SHA-256");
        byte[] key = sha.digest(password.getBytes("UTF-8"));
        key = Arrays.copyOf(key, 16); // AES-128
        return new SecretKeySpec(key, ALGORITHM);
    }

    /**
     * 处理文件/文件夹（加密）
     * @param path 文件或文件夹路径
     * @param password 密码
     */
    public static void encrypt(String path, String password) throws Exception {
        File file = new File(path);
        if (!file.exists()) {
            throw new FileNotFoundException("路径不存在: " + path);
        }

        if (file.isDirectory()) {
            // 处理文件夹（递归）
            File[] files = file.listFiles();
            if (files != null) {
                for (File f : files) {
                    encrypt(f.getAbsolutePath(), password);
                }
            }
        } else {
            // 处理单个文件（跳过已加密文件）
            if (file.getName().endsWith(ENCRYPT_FLAG)) {
                System.out.println("跳过已加密文件: " + file.getAbsolutePath());
                return;
            }
            encryptFile(file.getAbsolutePath(), password);
        }
    }

    /**
     * 处理文件/文件夹（解密）
     * @param path 文件或文件夹路径
     * @param password 密码
     */
    public static void decrypt(String path, String password) throws Exception {
        File file = new File(path);
        if (!file.exists()) {
            throw new FileNotFoundException("路径不存在: " + path);
        }

        if (file.isDirectory()) {
            // 处理文件夹（递归）
            File[] files = file.listFiles();
            if (files != null) {
                for (File f : files) {
                    decrypt(f.getAbsolutePath(), password);
                }
            }
        } else {
            // 处理单个文件（只解密带标记的文件）
            if (!file.getName().endsWith(ENCRYPT_FLAG)) {
                System.out.println("跳过非加密文件: " + file.getAbsolutePath());
                return;
            }
            decryptFile(file.getAbsolutePath(), password);
        }
    }

    /**
     * 加密单个文件（添加加密标记）
     */
    private static void encryptFile(String filePath, String password) throws Exception {
        File originalFile = new File(filePath);
        String encryptedPath = filePath + ENCRYPT_FLAG;

        // 加密到临时文件
        File tempFile = File.createTempFile("encrypt_temp", null);
        try (FileInputStream in = new FileInputStream(originalFile);
             FileOutputStream out = new FileOutputStream(tempFile);
             CipherOutputStream cipherOut = new CipherOutputStream(out, getEncryptCipher(password))) {

            byte[] buffer = new byte[8192];
            int bytesRead;
            while ((bytesRead = in.read(buffer)) != -1) {
                cipherOut.write(buffer, 0, bytesRead);
            }
        }

        // 替换原文件并添加加密标记
        replaceFile(tempFile.getAbsolutePath(), encryptedPath);
        // 删除原始文件
        if (!originalFile.delete()) {
            throw new IOException("无法删除原始文件: " + filePath);
        }
//        System.out.println("加密完成: " + encryptedPath);
    }

    /**
     * 解密单个文件（移除加密标记）
     */
    private static void decryptFile(String filePath, String password) throws Exception {
        if (!filePath.endsWith(ENCRYPT_FLAG)) {
            throw new IllegalArgumentException("不是加密文件: " + filePath);
        }

        File encryptedFile = new File(filePath);
        String originalPath = filePath.substring(0, filePath.length() - ENCRYPT_FLAG.length());

        // 解密到临时文件
        File tempFile = File.createTempFile("decrypt_temp", null);
        try (FileInputStream in = new FileInputStream(encryptedFile);
             CipherInputStream cipherIn = new CipherInputStream(in, getDecryptCipher(password));
             FileOutputStream out = new FileOutputStream(tempFile)) {

            byte[] buffer = new byte[8192];
            int bytesRead;
            while ((bytesRead = cipherIn.read(buffer)) != -1) {
                out.write(buffer, 0, bytesRead);
            }
        }

        // 替换加密文件并恢复原始文件名
        replaceFile(tempFile.getAbsolutePath(), originalPath);
        // 删除加密文件
        if (!encryptedFile.delete()) {
            throw new IOException("无法删除加密文件: " + filePath);
        }
//        System.out.println("解密完成: " + originalPath);
    }

    /**
     * 获取加密模式的Cipher
     */
    private static Cipher getEncryptCipher(String password) throws Exception {
        SecretKeySpec key = generateKey(password);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        return cipher;
    }

    /**
     * 获取解密模式的Cipher
     */
    private static Cipher getDecryptCipher(String password) throws Exception {
        SecretKeySpec key = generateKey(password);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, key);
        return cipher;
    }

    /**
     * 替换文件（确保操作原子性）
     */
    private static void replaceFile(String sourcePath, String targetPath) throws IOException {
        File sourceFile = new File(sourcePath);
        File targetFile = new File(targetPath);

        // 确保目标文件可写
        if (targetFile.exists() && !targetFile.canWrite()) {
            throw new IOException("目标文件不可写: " + targetPath);
        }

        // 删除目标文件（避免重名问题）
        if (targetFile.exists() && !targetFile.delete()) {
            throw new IOException("无法删除目标文件: " + targetPath);
        }

        // 移动临时文件到目标路径
        if (!sourceFile.renameTo(targetFile)) {
            throw new IOException("文件替换失败: " + targetPath);
        }
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        try {
            System.out.println("请选择操作：1-加密 2-解密");
            int choice = scanner.nextInt();
            scanner.nextLine(); // 消耗换行符

            System.out.println("请输入文件或文件夹路径：");
            String path = scanner.nextLine();
            File file = new File(path);
            if (!file.exists()) {
                System.out.println("路径不存在！");
                return;
            }

//            System.out.println("请输入密码：");
            String password = "258369" ;//scanner.nextLine();

            if (choice == 1) {
                System.out.println("开始加密...");
                encrypt(path, password);
                System.out.println("加密操作完成！");
            } else if (choice == 2) {
                System.out.println("开始解密...");
                decrypt(path, password);
                System.out.println("解密操作完成！");
            } else {
                System.out.println("无效选择");
            }
        } catch (Exception e) {
            System.err.println("操作失败：" + e.getMessage());
            e.printStackTrace();
        } finally {
            scanner.close();
        }
    }
}